package net.leesrobots.data.sink;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.google.common.collect.Lists;
import org.apache.commons.lang.StringUtils;
import org.apache.flume.*;
import org.apache.flume.conf.Configurable;
import org.apache.flume.sink.AbstractSink;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.UnsupportedEncodingException;
import java.sql.*;
import java.text.SimpleDateFormat;
import java.util.*;


/**
 * 数据库 Sink
 */
public class DBSink extends AbstractSink implements Configurable {
    private static final Logger logger = LoggerFactory.getLogger(DBSink.class);

    /**
     *
     */
    private String dbUrl;

    /**
     *
     */
    private String dbName;

    /**
     *
     */
    private String dbUser;

    /**
     *
     */
    private String dbPass;

    /**
     *
     */
    private String dbTable;

    /**
     *
     */
    private String dbTableSQL;

    /**
     *
     */
    private String dbTableTimeformat;

    /**
     * 是否自动创建表
     */
    private Boolean dbTableAutocreate;

    /**
     * 每次多少条记录写入一次数据库
     */
    private int batchSize;

//    private DataSource dataSource;

    /**
     *
     */
    private CounterGroup counterGroup;

    /**
     * 启动时间
     */
    private long startTime;

    private HashMap<String, String> tableMetas;

    /**
     * 构造函数
     */
    public DBSink() {
        this.counterGroup = new CounterGroup();
        counterGroup.setName("Database counter");

        startTime = System.currentTimeMillis();
    }

    /**
     * 获取配置
     * @param context Context
     */
    @Override
    public void configure(Context context) {
        dbUrl = context.getString("db.url");
        dbName = context.getString("db.name");
        dbUser = context.getString("db.user");
        dbPass = context.getString("db.pass");
        dbTable = context.getString("db.table");
        dbTableSQL = context.getString("db.table.sql");
        dbTableTimeformat = context.getString("db.table.timeformat", "yyyy-MM-dd");
        dbTableAutocreate = context.getBoolean("db.table.autocreate", true);
        batchSize = Integer.parseInt(context.getString("batch_size", "100"));

        logger.info("=============================================================");
        logger.info("url: {}, name: {}", dbUrl, dbUser);
        logger.info("pass: {}, table: {}", dbPass, getTable());
        logger.info("=============================================================");
    }

    @Override
    public synchronized void start() {
        logger.info("Database sink {} starting", getName());
        tableMetas = new HashMap<String, String>();

        try {
            Class.forName("com.mysql.jdbc.Driver");
            Connection connection = DriverManager.getConnection(dbUrl, dbUser, dbPass);

            checkTableExists(connection);

            DatabaseMetaData metaData = connection.getMetaData();
            ResultSet colRet = metaData.getColumns(null, "%", getTable(), "%");

            while(colRet.next()) {
                String columnName = colRet.getString("COLUMN_NAME");
                String columnType = colRet.getString("TYPE_NAME");
                tableMetas.put(columnName, columnType);
                logger.info(">>>>>>>>>>>>>>>>>> {} - {}", columnName, columnType);
            }

            colRet.close();
            connection.close();
        } catch (SQLException e) {
            logger.error(e.getMessage());
        } catch (ClassNotFoundException e) {
            logger.error(e.getMessage());
        }
        super.start();

        logger.info("Database sink {} started", this.getName());
    }

    private void checkTableExists(Connection connection) throws SQLException {
        if (dbTableAutocreate) {
            logger.warn("check table {} start", getTable());
            String sql = String.format("SELECT count(0) as C FROM information_schema.TABLES WHERE table_schema='%s' AND TABLE_NAME='%s'", dbName, getTable());
            Statement statement = connection.createStatement(ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_UPDATABLE);
            ResultSet rs = statement.executeQuery(sql);
            rs.next();
            if (rs.getInt("C") == 0) {
                logger.warn("Create table {} begin", getTable());
                logger.error(getTableSQL());
                PreparedStatement pst = connection.prepareStatement(getTableSQL());
                pst.executeUpdate();
                pst.close();
                logger.warn("Create table {} end", getTable());
            }
            rs.close();
            statement.close();
        }
    }

    @Override
    public synchronized void stop() {
        logger.info("Database sink {} stopping", getName());
        super.stop();

        logger.info("Database sink {} stopped. Metrics: {}", this.getName(), counterGroup);
    }

    /**
     * 返回数据表名
     * @return 数据表名
     */
    private String getTable() {
        SimpleDateFormat simpleDateFormat = new SimpleDateFormat(dbTableTimeformat);
        String formattedTime = simpleDateFormat.format(new java.util.Date());

        return String.format(dbTable, formattedTime);
    }

    private String getTableSQL() {
        SimpleDateFormat simpleDateFormat = new SimpleDateFormat(dbTableTimeformat);
        String formattedTime = simpleDateFormat.format(new java.util.Date());

        return String.format(dbTableSQL, formattedTime);
    }

    /**
     * 处理一组事件
     * @param events 事件组
     */
    private void processEvents(List<Event> events) throws SQLException, ClassNotFoundException {
        Connection connection = null;
        try {

            connection = DriverManager.getConnection(dbUrl, dbUser, dbPass);
            checkTableExists(connection);
            connection.setAutoCommit(false);

            List<String> fields = new ArrayList<String>();
            List<String> placeholders = new ArrayList<String>();
            for(String field : this.tableMetas.keySet()) {
                fields.add(field);
                placeholders.add("?");
            }

            String sql;
            sql = String.format("REPLACE INTO `%s`(%s) VALUES(%s)", getTable(), StringUtils.join(fields, ","), StringUtils.join(placeholders, ","));
            PreparedStatement ps = connection.prepareStatement(sql);
            String ds;
            byte[] datas;
            for (Event event : events) {
                try {
                    datas = event.getBody();
                    ds = new String(datas, "UTF-8");
                    JSONObject object = JSON.parseObject(ds);
                    int i = 1;
                    for (String field: this.tableMetas.keySet()) {
                        ps.setString(i++, object.getString(field));
                    }
                    ps.addBatch();
                } catch (UnsupportedEncodingException e) {
                    logger.error("UnsupportedEncodingException: ", e);
                }
            }
            ps.executeBatch();
            connection.commit();
        } catch (SQLException e) {
            logger.error("SQL Exception: ", e);
            if (connection != null) {
                connection.rollback();
            }
            throw e;
        } finally {
            if (connection != null && !connection.isClosed()) {
                connection.close();
            }
        }
    }

    /**
     * 处理事件
     * @return Status
     * @throws EventDeliveryException EventDeliveryException
     */
    @Override
    public Status process() throws EventDeliveryException {
        Status status = Status.READY;
        Channel channel = getChannel();
        Transaction transaction = channel.getTransaction();

        transaction.begin();
        try {
            logger.info("Channel {} transaction open", channel.getName());

            List<Event> batch = Lists.newLinkedList();
            for (int i = 0; i < batchSize; i++) {
                Event event = channel.take();
                logger.debug("Event {}", counterGroup.get("executeSize"));

                if (event == null) {
                    break;
                }

                counterGroup.incrementAndGet("executeSize");
                batch.add(event);
            }

            int size = batch.size();

            if (size == 0) {
                status = Status.BACKOFF;
            } else {
                processEvents(batch);
            }

            transaction.commit();
            logger.info("Database sink {} running. Metrics: {}", getName(), counterGroup);
        } catch (Throwable  e) {
            transaction.rollback();
            logger.error(getName() + ": Unable to get event from channel. Exception follows.", e);
            status = Status.BACKOFF;
            if (e instanceof Error) {
                throw (Error)e;
            }
        } finally {
            transaction.close();
        }

        return status;
    }
}